{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "038c9d4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp wer_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb2c63ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c50ffe3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "import jiwer\n",
    "from whisper_normalizer.english import EnglishTextNormalizer\n",
    "\n",
    "import torchaudio\n",
    "from pathlib import Path\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1606cd1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "engnorm = EnglishTextNormalizer()\n",
    "def whisper_normalize(x):\n",
    "    if type(x) == list:\n",
    "        return [engnorm(y) for y in x]\n",
    "    else:\n",
    "        return engnorm(x)\n",
    "\n",
    "default_transform = jiwer.transforms.Compose([\n",
    "    jiwer.transforms.ToLowerCase(),\n",
    "    jiwer.transforms.ExpandCommonEnglishContractions(),\n",
    "    whisper_normalize,\n",
    "    jiwer.transforms.RemoveMultipleSpaces(),\n",
    "    jiwer.transforms.Strip(),\n",
    "    jiwer.transforms.RemovePunctuation(),\n",
    "    jiwer.transforms.ReduceToListOfListOfWords(),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c85da4f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['footnote', 'somber', 'tashan', 'may', '12', '1856'],\n",
       " ['footnote', 'sumner', 'to', 'shannon', 'may', '12th', '1856']]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "default_transform([\"Footnote, Somber Tashan, May 12, 1856\", \"FOOTNOTE SUMNER TO SHANNON MAY TWELFTH EIGHTEEN FIFTY SIX\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8af761b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def librispeech_data(datadir, sample_rate=16000):\n",
    "    for file in Path(datadir).rglob('*.txt'):\n",
    "        for line in file.read_text().split('\\n'):\n",
    "            if not line: continue\n",
    "            idx, text = line.split(\" \", 1)\n",
    "            x, sr = torchaudio.load((file.parent/idx).with_suffix('.flac'))\n",
    "            if sr != sample_rate:\n",
    "                x = torchaudio.transforms.Resample(sr, self.sample_rate)(x)\n",
    "            yield x, text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d76a8b39",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class DfBuilder:\n",
    "    def __init__(self):\n",
    "        self.data = {}\n",
    "        \n",
    "    def push(self, **kwargs):\n",
    "        for k,v in kwargs.items():\n",
    "            if k not in self.data:\n",
    "                self.data[k] = [v]\n",
    "            else:\n",
    "                self.data[k].append(v)\n",
    "    \n",
    "    def df(self):\n",
    "        return pd.DataFrame(self.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57959900",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class WERStats(DfBuilder):\n",
    "    def __init__(self, transform=default_transform):\n",
    "        super().__init__()\n",
    "        self.reference_transform = transform\n",
    "        self.hypothesis_transform = transform\n",
    "    \n",
    "    def push_sample(self, snd, gt_text, text, idx=None):\n",
    "        if snd is not None: self.push(secs = snd.shape[-1]/16000)\n",
    "        diff = jiwer.process_words(gt_text, text, reference_transform=self.reference_transform, hypothesis_transform=self.hypothesis_transform)\n",
    "        self.push(\n",
    "            idx = idx,\n",
    "            gt_text = gt_text,\n",
    "            text = text,\n",
    "            wer = diff.wer,\n",
    "            mer = diff.mer,\n",
    "            wil = diff.wil,\n",
    "            wip = diff.wip,\n",
    "        )\n",
    "        return diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bb37313",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/nbdev/export.py:54: UserWarning: Notebook '/root/workspace/spear-tts-pytorch/nbs/2B. VQ semantic token extraction evaluation.ipynb' uses `#|export` without `#|default_exp` cell.\n",
      "Note nbdev2 no longer supports nbdev1 syntax. Run `nbdev_migrate` to upgrade.\n",
      "See https://nbdev.fast.ai/getting_started.html for more information.\n",
      "  warn(f\"Notebook '{nbname}' uses `#|export` without `#|default_exp` cell.\\n\"\n",
      "/opt/conda/lib/python3.10/site-packages/nbdev/export.py:54: UserWarning: Notebook '/root/workspace/spear-tts-pytorch/nbs/2F. Residual (RQ) semantic token extraction model - dataloading performance tests.ipynb' uses `#|export` without `#|default_exp` cell.\n",
      "Note nbdev2 no longer supports nbdev1 syntax. Run `nbdev_migrate` to upgrade.\n",
      "See https://nbdev.fast.ai/getting_started.html for more information.\n",
      "  warn(f\"Notebook '{nbname}' uses `#|export` without `#|default_exp` cell.\\n\"\n",
      "/opt/conda/lib/python3.10/site-packages/nbdev/export.py:54: UserWarning: Notebook '/root/workspace/spear-tts-pytorch/nbs/2F. Residual (RQ) semantic token extraction model helpers.ipynb' uses `#|export` without `#|default_exp` cell.\n",
      "Note nbdev2 no longer supports nbdev1 syntax. Run `nbdev_migrate` to upgrade.\n",
      "See https://nbdev.fast.ai/getting_started.html for more information.\n",
      "  warn(f\"Notebook '{nbname}' uses `#|export` without `#|default_exp` cell.\\n\"\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe1b4889",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
